Lesson 2: Image Segmentation

⏳ Note (Kernel Starting): This notebook takes about 30 seconds to be ready to use. You may start and watch the video while you wait.

• In this classroom, the libraries have been already installed for you.
• If you would like to run this code on your own machine, you need to install the following:

  !pip install ultralytics torch
  

Load the sample image

from PIL import Image
raw_image = Image.open("dogs.jpg")
raw_image

Note: the images referenced in this notebook have already been uploaded to the Jupyter directory, in this classroom, for your convenience. For further details, please refer to the Appendix section located at the end of the lessons.

• Resize the image.

from utils import resize_image
resized_image = resize_image(raw_image, input_size=1024)
resized_image

Import and prepare the model

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device

device(type='cpu')

Info about torch.

from ultralytics import YOLO
model = YOLO('./FastSAM.pt')

Info about 'FastSAM'

Use the model

Note: utils is an additional file containing the methods that have been already developed for you to be used in this classroom. For further details, please refer to the Appendix section located at the end of the lessons.

from utils import show_points_on_image

# Define the coordinates for the point in the image
# [x_axis, y_axis]
input_points = [ [350, 450 ] ]

input_labels = [1] # positive point

# Function written in the utils file
show_points_on_image(resized_image, input_points)

# Run the model
results = model(resized_image, device=device, retina_masks=True)

0: 704x1024 23 objects, 7193.8ms
Speed: 3.5ms preprocess, 7193.8ms inference, 195.1ms postprocess per image at shape (1, 3, 704, 1024)

• Filter the mask based on the point defined before.

from utils import format_results, point_prompt

results = format_results(results[0], 0)

# Generate the masks
masks, _ = point_prompt(results, input_points, input_labels)

from utils import show_masks_on_image

# Visualize the generated masks
show_masks_on_image(resized_image, [masks])

• Define 'semantic masks' - two points to be masked.

# Specify two points in the same image
# [x_axis, y_axis]
input_points = [ [350, 450], [620, 450] ]

# Specify both points as "positive prompt"
input_labels = [1 , 1] # both positive points

# Visualize the points defined before
show_points_on_image(resized_image, input_points)

# Run the model
results = model(resized_image, device=device, retina_masks=True)

0: 704x1024 23 objects, 6979.6ms
Speed: 2.6ms preprocess, 6979.6ms inference, 197.8ms postprocess per image at shape (1, 3, 704, 1024)

results = format_results(results[0], 0)

# Generate the masks
masks, _ = point_prompt(results, input_points, input_labels)

# Visualize the generated masks
show_masks_on_image(resized_image, [masks])

Note: Please note that the results obtained from running this notebook may vary slightly from those demonstrated by the instructor in the video.

• Identify subsections of the image by adding a negative prompt.

# Define the coordinates for the regions to be masked
# [x_axis, y_axis]
input_points = [ [350, 450], [400, 300]  ]

input_labels = [1, 0] # positive prompt, negative prompt

# Visualize the points defined above
show_points_on_image(resized_image, input_points, input_labels)

Note: From the image above, the red star indicates the negative prompt and the green star the positive prompt.

# Run the model
results = model(resized_image, device=device, retina_masks=True)

0: 704x1024 23 objects, 6023.0ms
Speed: 2.7ms preprocess, 6023.0ms inference, 201.0ms postprocess per image at shape (1, 3, 704, 1024)

results = format_results(results[0], 0)

# Generate the masks
masks, _ = point_prompt(results, input_points, input_labels)

# Visualize the generated masks
show_masks_on_image(resized_image, [masks])

Note: From the image above, only the jacket, from the dog in the left, was segmented, so, it is following the command given by the positive prompt!

Prompting with bounding boxes

from utils import box_prompt

# Set the bounding box coordinates
# [xmin, ymin, xmax, ymax]
input_boxes = [530, 180, 780, 600]

from utils import show_boxes_on_image

# Visualize the bounding box defined with the coordinates above
show_boxes_on_image(resized_image, [input_boxes])

• Now, try to isolate the mask from the total output of the model.

from utils import box_prompt

results = model(resized_image, device=device, retina_masks=True)

0: 704x1024 23 objects, 6111.7ms
Speed: 2.6ms preprocess, 6111.7ms inference, 289.9ms postprocess per image at shape (1, 3, 704, 1024)

# Generate the masks
masks = results[0].masks.data

masks

tensor([[[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]],

        [[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]],

        [[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]],

        ...,

        [[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]],

        [[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]],

        [[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]]])

# Convert to True/False boolean mask
masks = masks > 0

masks

tensor([[[False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         ...,
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False]],

        [[False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         ...,
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False]],

        [[False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         ...,
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False]],

        ...,

        [[False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         ...,
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False]],

        [[False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         ...,
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False]],

        [[False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         ...,
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False],
         [False, False, False,  ..., False, False, False]]])

masks, _ = box_prompt(masks, input_boxes)

# Visualize the masks
show_masks_on_image(resized_image, [masks])

# Print the segmentation mask, but in its raw format
masks

array([[False, False, False, ..., False, False, False],
       [False, False, False, ..., False, False, False],
       [False, False, False, ..., False, False, False],
       ...,
       [False, False, False, ..., False, False, False],
       [False, False, False, ..., False, False, False],
       [False, False, False, ..., False, False, False]])

# To visualize, import matplotlib
from matplotlib import pyplot as plt

# Plot the binary mask as an image
plt.imshow(masks, cmap='gray')

<matplotlib.image.AxesImage at 0x7f6014914710>

Try yourself!

Try the image segmentation explained before with your own images.

# To start opening images, upload your own or use the sample images we've uploaded, for example: younes.png
# The image younes.png is already uploaded in this classroom
raw_image = Image.open('younes.png')
raw_image

# Resize image
from utils import resize_image
resized_image = resize_image(raw_image, input_size=1024)
resized_image

# Define the coordinates for the point: [x_axis, y_axis]
input_points = [ [700, 450 ] ]

# Define the positive or negative prompt
input_labels = [1] 

show_points_on_image(resized_image, input_points)

# Run the model
results = model(resized_image, device=device, retina_masks=True)
results = format_results(results[0], 0)
# Generate the masks
masks, _ = point_prompt(results, input_points, input_labels)
# Visualize the generated masks
show_masks_on_image(resized_image, [masks])

0: 640x1024 99 objects, 7412.4ms
Speed: 3.3ms preprocess, 7412.4ms inference, 993.9ms postprocess per image at shape (1, 3, 640, 1024)

Additional Resources

• For more on how to use Comet for experiment tracking, check out this Quickstart Guide and the Comet Docs.
• This course was based off a set of two blog articles from Comet. Explore them here for more on how to use newer versions of Stable Diffusion in this pipeline, additional tricks to improve your inpainting results, and a breakdown of the pipeline architecture:
  • SAM + Stable Diffusion for Text-to-Image Inpainting
  • Image Inpainting for SDXL 1.0 Base Model + Refiner